import numpy as np
import matplotlib.pyplot as plt
import os
from matplotlib.colors import LogNorm

fname = 'figS1.npz'

# load npz file
data = np.load(fname, allow_pickle=True)
print(data.files)
n = data['n']
B = data['B']
Rxx_sym = data['Rxx_sym']
Rxx_raw = data['Rxx_raw']
Rxy_antisym = data['Rxy_antisym']
Rxy_raw = data['Rxy_raw']
Vtg = data['Vtg']
Vbg = data['Vbg']
n0 = data['n0'] # density at charge neutrality point
ns = data['ns'] # density at 4 electron per moire unit cell
device_name = data['device_name']
back_gate_thickness_nm = data['back_gate_thickness_nm']
top_gate_thickness_nm = data['top_gate_thickness_nm']

# plot
fig, ax = plt.subplots(1, 2, figsize=(10, 4))
fig.suptitle('Fig. figS1 raw data')
im = ax[0].pcolormesh(n, B, Rxx_raw, cmap='inferno',)
ax[0].set_xlabel('Density ($10^{12}$ cm$^{-2}$)')
ax[0].set_ylabel('$B$ (T)')
cbar = fig.colorbar(im, ax=ax[0])
cbar.set_label('$R_{xx}$ raw ($\Omega$)')

im = ax[1].pcolormesh(n, B, Rxy_raw, cmap='seismic',)
ax[1].set_xlabel('Density ($10^{12}$ cm$^{-2}$)')
ax[1].set_ylabel('$B$ (T)')
cbar = fig.colorbar(im, ax=ax[1])
cbar.set_label('$R_{xy}$ raw ($\Omega$)')
plt.show()


# plot
fig, ax = plt.subplots(1, 2, figsize=(10, 4))
fig.suptitle('figS1 symmetrized/antisymmetrized data')
im = ax[0].pcolormesh(n, B, Rxx_sym, cmap='inferno', norm=LogNorm(vmin=1e2, vmax=1e4))
ax[0].set_xlabel('Density ($10^{12}$ cm$^{-2}$)')
ax[0].set_ylabel('$B$ (T)')
cbar = fig.colorbar(im, ax=ax[0])
cbar.set_label('$R_{xx}$ symmetrized ($\Omega$)')

im = ax[1].pcolormesh(n, B, Rxy_antisym, cmap='seismic',)
ax[1].set_xlabel('Density ($10^{12}$ cm$^{-2}$)')
ax[1].set_ylabel('$B$ (T)')
cbar = fig.colorbar(im, ax=ax[1])
cbar.set_label('$R_{xy}$ antisymmetrized ($\Omega$)')
plt.show()